%Figure 2
clear variables

load('Fig2_S2_S3_mobility.mat')

%% Figure 2 Panel A

figure
ax = usamap('all');
set(ax, 'Visible', 'off')
states = shaperead('usastatelo', 'UseGeoCoords', true);
names = {states.Name};
indexHawaii = strcmp('Hawaii',names);
indexAlaska = strcmp('Alaska',names);
indexConus = 1:numel(states);
indexConus(indexHawaii|indexAlaska) = []; 

for i=1:numel(states)-1 %fills everything leaving DC for last
    if i<9
        states(i).a=log10(a_new(i));
    else
        states(i).a=log10(a_new(i+1));
    end
end

states(51).a=log10(a_new(9));

% log scale
surfaceColors = makesymbolspec('Polygon', {'a', ...
    [-3 log10(a_new(31))], 'FaceColor', flipud(hot(12)) });

% Display the three regions.
geoshow(ax(1), states(indexConus),  'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')
geoshow(ax(2), states(indexAlaska), 'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')
geoshow(ax(3), states(indexHawaii), 'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')

% Hide the frame.
caxis([-3 log10(a_new(31))])
colormap(flipud(hot(12)))

colorbar
for k = 1:3
    setm(ax(k), 'Frame', 'off', 'Grid', 'off',...
      'ParallelLabel', 'off', 'MeridianLabel', 'off')
end

%% Figure 2 Panel B

figure
ax = usamap('all');
set(ax, 'Visible', 'off')
states = shaperead('usastatelo', 'UseGeoCoords', true);
names = {states.Name};
indexHawaii = strcmp('Hawaii',names);
indexAlaska = strcmp('Alaska',names);
indexConus = 1:numel(states);
indexConus(indexHawaii|indexAlaska) = []; 

for i=1:numel(states)-1 %fills everything leaving DC for last
    if i<9
        states(i).b=b_new(i);
    else
        states(i).b=b_new(i+1);
    end
end

states(51).b=b_new(9);

surfaceColors = makesymbolspec('Polygon', {'b', ...
    [1.5 max(b_new)], 'FaceColor', flipud(hot(12)) });

% Display the three regions.

geoshow(ax(1), states(indexConus),  'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')
geoshow(ax(2), states(indexAlaska), 'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')
geoshow(ax(3), states(indexHawaii), 'DisplayType','polygon','SymbolSpec', surfaceColors,'DefaultFaceColor','white')
% Hide the frame.
caxis([1.5 max(b_new)])
colormap(flipud(hot(12)))

colorbar
for k = 1:3
    setm(ax(k), 'Frame', 'off', 'Grid', 'off',...
      'ParallelLabel', 'off', 'MeridianLabel', 'off')
end

%% Figure 2 Panel C
figure
scatter(log10(table2array(county_unacast(1:end-1,5))),log10(county_fit(:,1)))
xlabel('Log of population density')
ylabel('Log of a')
text(3.5,3.5,'R{^2}=0.84','FontSize',14)

%% Figure 2 Panel D
figure
scatter(log10(table2array(county_unacast(1:end-1,5))),county_fit(:,2))
xlabel('Log of population density')
ylabel('b')
text(3.5,5,'R{^2}=0.15','FontSize',14)
